/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020. All rights reserved.
 * secGear is licensed under the Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
 * PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

#include <stdio.h>
#include <unistd.h>
#include <linux/limits.h>
#include "enclave.h"
#include "hotcalls_u.h"
#include "hotcalls.h"
#include "common.h"
#include "string.h"

# define MAX_PATH FILENAME_MAX

cc_enclave_t globalEnclaveID;
cc_enclave_result_t res = CC_FAIL;

typedef cc_enclave_result_t (*EcallFunction)(cc_enclave_t, void* );
typedef uint8_t cc_enclave_launch_t[1024];

#define PERFORMANCE_MEASUREMENT_NUM_REPEATS 10000
#define MEASUREMENTS_ROOT_DIR               "measurments"

using namespace std;
/*
cpu时钟周期数
*/
inline __attribute__((always_inline))  uint64_t rdtscp(void)
{
        unsigned int low, high;

        asm volatile("rdtscp" : "=a" (low), "=d" (high));

        return low | ((uint64_t)high) << 32;
}

/* OCall functions */
void ocall_print_string(const char *str)
{
    /* Proxy/Bridge will check the length and null-terminate 
     * the input string to prevent buffer overflow. 
     */
    printf("%s", str);
}

void* EnclaveResponderThread( void* hotEcallAsVoidP )
{
    //To be started in a new thread
    HotCall *hotEcall = (HotCall*)hotEcallAsVoidP;
    EcallStartResponder( globalEnclaveID, hotEcall );

    return NULL;
}

void MyCustomOcall( void* data )
{
    //Because RDTSCP is not allowed inside an enclave in SGX 1.x, we have to issue it here,
    //in the ocall. Therefore, instead of measuring enclave-->ocall-->enclave, we will measure 
    //ocall-->enclave-->next_ocall
    
    static uint64_t startTime     = 0;  

    OcallParams* ocallParams = (OcallParams*)data;

    *(ocallParams->cyclesCount)  = rdtscp() - startTime; //startTime was set in previous iteration (except when first called)
    ocallParams->counter++;

    startTime     = rdtscp(); //for next iteration
}

void* OcallResponderThread( void* hotCallAsVoidP )
{
    void (*callbacks[1])(void*);
    callbacks[0] = MyCustomOcall;

    HotCallTable callTable;
    callTable.numEntries = 1;
    callTable.callbacks  = callbacks;

    HotCall_waitForCall( (HotCall *)hotCallAsVoidP, &callTable );

    return NULL;
}

class HotCallsTester {
public:
    HotCallsTester() {
        m_enclaveID = 0;

        if( initialize_enclave() < 0){
            printf("Enter a character before exit ...\n");
            getchar();
            throw HotCallsTesterError(); 
        }

        CreateMeasurementsDirectory();
    }

    ~HotCallsTester() {
        /* Destroy the enclave */
        sgx_destroy_enclave( m_enclaveID );
    }

    void Run( void ) {
        TestHotEcalls();
        TestHotOcalls();

        TestSDKEcalls();
        TestSDKOcalls();
    }

    void TestHotEcalls()
    {
        uint64_t performaceMeasurements[ PERFORMANCE_MEASUREMENT_NUM_REPEATS ]= {0};

        uint64_t    startTime       = 0;
        uint64_t    endTime         = 0;
        int         data            = 0;
        int         expectedData    = 0;
        HotCall     hotEcall        = HOTCALL_INITIALIZER;
        hotEcall.data               = &data; 

        globalEnclaveID = m_enclaveID;
        pthread_create(&hotEcall.responderThread, NULL, EnclaveResponderThread, (void*)&hotEcall);

        const uint16_t requestedCallID = 0;
        for( uint64_t i=0; i < PERFORMANCE_MEASUREMENT_NUM_REPEATS; ++i ) {
            startTime = rdtscp();
            HotCall_requestCall( &hotEcall, requestedCallID, &data );
            endTime   = rdtscp();
        
            performaceMeasurements[ i ] = endTime       - startTime;

            expectedData++;
            if( data != expectedData ){
                printf( "Error! Data is different than expected: %d != %d\n", data, expectedData );
            }
        }

        StopResponder( &hotEcall );
        ostringstream filename;
        filename <<  "HotEcall_latencies_in_cycles.csv";
        WriteMeasurementsToFile( filename.str(), 
                                 (uint64_t*)performaceMeasurements, 
                                 PERFORMANCE_MEASUREMENT_NUM_REPEATS ) ;
    }

    void TestSDKEcalls()
    {
        uint64_t performaceMeasurements[ PERFORMANCE_MEASUREMENT_NUM_REPEATS ]= {0};

        uint64_t    startTime       = 0;
        uint64_t    endTime         = 0;
        int         data            = 0;
        int         expectedData    = 0;

        globalEnclaveID = m_enclaveID;        

        const uint16_t requestedCallID = 0;
        for( uint64_t i=0; i < PERFORMANCE_MEASUREMENT_NUM_REPEATS; ++i ) {
            startTime = rdtscp();
            MyCustomEcall( m_enclaveID, &data );
            endTime   = rdtscp();
        
            performaceMeasurements[ i ] = endTime       - startTime;

            expectedData++;
            if( data != expectedData ){
                printf( "Error! Data is different than expected: %d != %d\n", data, expectedData );
            }
        }

        ostringstream filename;
        filename <<  "SDKEcall_latencies_in_cycles.csv";
        WriteMeasurementsToFile( filename.str(), 
                                 (uint64_t*)performaceMeasurements, 
                                 PERFORMANCE_MEASUREMENT_NUM_REPEATS ) ;
    }

    void TestHotOcalls()
    {
        uint64_t performaceMeasurements[ PERFORMANCE_MEASUREMENT_NUM_REPEATS ]= {0};

        OcallParams ocallParams;
        ocallParams.counter     = 0;
        HotCall     hotOcall    = HOTCALL_INITIALIZER;
        hotOcall.data           = &ocallParams;
        
        pthread_create( &hotOcall.responderThread, NULL, OcallResponderThread, (void*)&hotOcall );
       
        EcallMeasureHotOcallsPerformance( 
                m_enclaveID, 
                (uint64_t*)performaceMeasurements, 
                PERFORMANCE_MEASUREMENT_NUM_REPEATS,
                &hotOcall );
        StopResponder( &hotOcall );

        ostringstream filename;
        filename <<  "HotOcall_latencies_in_cycles.csv";
        WriteMeasurementsToFile( filename.str(), 
                                 (uint64_t*)performaceMeasurements, 
                                 PERFORMANCE_MEASUREMENT_NUM_REPEATS ) ;
    }

    void TestSDKOcalls()
    {
        uint64_t performaceMeasurements[ PERFORMANCE_MEASUREMENT_NUM_REPEATS ]= {0};

        OcallParams ocallParams;
        ocallParams.counter     = 0;
        HotCall     hotOcall    = HOTCALL_INITIALIZER;
        hotOcall.data           = &ocallParams;
        
        EcallMeasureSDKOcallsPerformance( 
                m_enclaveID, 
                (uint64_t*)performaceMeasurements, 
                PERFORMANCE_MEASUREMENT_NUM_REPEATS,
                &ocallParams );
        
        ostringstream filename;
        filename <<  "SDKOcall_latencies_in_cycles.csv";
        WriteMeasurementsToFile( filename.str(), 
                                 (uint64_t*)performaceMeasurements, 
                                 PERFORMANCE_MEASUREMENT_NUM_REPEATS ) ;
    }

private:
    /* Global EID shared by multiple threads */
    cc_enclave_t m_enclaveID;

    int              m_sgxDriver;
    string           m_measurementsDir;

    void WriteMeasurementsToFile( string fileName, uint64_t* measurementsMatrix, size_t numRows )
    {
        string fileFullPath = m_measurementsDir + "/" + fileName;
        cout << "Writing results.. ";
        cout << fileFullPath << " ";
        ofstream measurementsFile;
        measurementsFile.open( fileFullPath, ios::app );
        for( size_t rowIdx = 0; rowIdx < numRows; ++rowIdx ) {
            measurementsFile << measurementsMatrix[ rowIdx ] << " ";
            measurementsFile << "\n";
        }
        
        measurementsFile.close();

        cout << "Done\n";
    }

    /* Initialize the enclave:
     *   Step 1: try to retrieve the launch token saved by last transaction
     *   Step 2: call sgx_create_enclave to initialize an enclave instance
     *   Step 3: save the launch token if it is updated
     */
    int initialize_enclave(void)
    {
        char *path = PATH;
        char real_p[MAX_PATH];
        char token_path[MAX_PATH] = {'\0'};
        cc_enclave_launch_t token = {0};
        cc_enclave_result_t ret = CC_ERROR_UNEXPECTED;
        int updated = 0;
        
        /* Step 1: try to retrieve the launch token saved by last transaction 
         *         if there is no token, then create a new one.
         */
        /* try to get the token saved in $HOME */
        const char *home_dir = getpwuid(getuid())->pw_dir;
        if (home_dir != NULL && 
            (strlen(home_dir)+strlen("/")+sizeof(TOKEN_FILENAME)+1) <= MAX_PATH) {
            /* compose the token path */
            strncpy(token_path, home_dir, strlen(home_dir));
            strncat(token_path, "/", strlen("/"));
            strncat(token_path, TOKEN_FILENAME, sizeof(TOKEN_FILENAME)+1);
        } else {
            /* if token path is too long or $HOME is NULL */
            strncpy(token_path, TOKEN_FILENAME, sizeof(TOKEN_FILENAME));
        }
        
        FILE *fp = fopen(token_path, "rb");
        if (fp == NULL && (fp = fopen(token_path, "wb")) == NULL) {
            printf("Warning: Failed to create/open the launch token file \"%s\".\n", token_path);
        }
        
        if (fp != NULL) {
            /* read the token from saved file */
            size_t read_num = fread(token, 1, sizeof(cc_enclave_launch_t), fp);
            if (read_num != 0 && read_num != sizeof(cc_enclave_launch_t)) {
                /* if token is invalid, clear the buffer */
                memset(&token, 0x0, sizeof(cc_enclave_t));
                printf("Warning: Invalid launch token read from \"%s\".\n", token_path);
            }
        }

        /* check file exists, if not exist then use absolute path */
        if (realpath(path, real_p) == NULL) {
            if (getcwd(real_p, sizeof(real_p)) == NULL) {
                printf("Cannot find enclave.sign.so");
                return -1;
            }
            if (PATH_MAX - strlen(real_p) <= strlen("/enclave.signed.so")) {
                printf("Failed to strcat enclave.sign.so path");
                return -1;
            }
            (void)strcat(real_p, "/enclave.signed.so");
        }


        /* Step 2: call sgx_create_enclave to initialize an enclave instance */
        /* Debug Support: set 2nd parameter to 1 */
        res = cc_enclave_create(real_p, AUTO_ENCLAVE_TYPE, 0, SECGEAR_DEBUG_FLAG, NULL, 0, &m_enclaveID);
        if (res != CC_SUCCESS) {
            printf("host create enclave error\n");
            return -1;
        }
        printf("host create enclave success\n");
        
        /* Step 3: save the launch token if it is updated */
        if (updated == FALSE || fp == NULL) {
            /* if the token is not updated, or file handler is invalid, do not perform saving */
            if (fp != NULL) fclose(fp);
            return 0;
        }
        
        /* reopen the file with write capablity */
        fp = freopen(token_path, "wb", fp);
        if (fp == NULL) return 0;
        size_t write_num = fwrite(token, 1, sizeof(cc_enclave_launch_t), fp);
        if (write_num != sizeof(cc_enclave_t))
            printf("Warning: Failed to save launch token to \"%s\".\n", token_path);
        fclose(fp);
        printf("line: %d\n", __LINE__ );
        return 0;
    }

    void CreateMeasurementsDirectory()
    {
        char timestamp[ 100 ] = {0};
        GetTimeStamp( timestamp, 100 );
        // printf( "%s\n", timestamp);

        if( ! IsDirectoryExists( MEASUREMENTS_ROOT_DIR )  ) {
            printf( "Creating directory %s\n", MEASUREMENTS_ROOT_DIR );
            mkdir(MEASUREMENTS_ROOT_DIR, 0700);
        }

        m_measurementsDir = string( MEASUREMENTS_ROOT_DIR ) + "/" + timestamp;
        if( ! IsDirectoryExists( m_measurementsDir )  ) {
            printf( "Creating directory %s\n", m_measurementsDir.c_str() );
            mkdir(m_measurementsDir.c_str(), 0700);
        }
    }

    void GetTimeStamp( char *timestamp, size_t size )
    {
      time_t rawtime;
      struct tm * timeinfo;
      time (&rawtime);
      timeinfo = localtime(&rawtime);

      strftime(timestamp,size,"%Y-%m-%d_%H-%M-%S",timeinfo);
    }

    bool IsDirectoryExists( string path )
    {
        struct stat st = {0};

        return ! (stat(path.c_str(), &st) == -1);
    }
};
/*
应用程序入口
*/

int main(int argc, char *argv[])
{
    (void)(argc);
    (void)(argv);

    HotCallsTester hotCallsTester;
    hotCallsTester.Run();

    return 0;
}
